"""De-identify all dicom images in a folder
Usage:
    batch_deidentification.py <input_dir> <output_dir> <except_dir>
"""  # 支持命令行输入
import os  # 操纵文件处理
import time
import zipfile  # 压缩包处理
import pydicom  # cdm图像处理
import shutil  # 文件删除
from tqdm import tqdm  # 进度条显示
from docopt import docopt  # 支持命令行输入
from de_id_logger import logger  # 导入本地日志处理类对象
from settigs_config import desensitization_switch  # 配置文件 是否修改
from settigs_config import desensitization_tags  # 配置文件，修改的值


class BatchDeId(object):

    def __init__(self, input_dir, output_dir, outs_dir):
        self.input_dir = input_dir  # 原件压缩包的读取路径
        self.output_dir = output_dir  # 存储修改后的dcm压缩包
        self.except_dir = "{}/{}".format(outs_dir, 'except_out')  # 处理异常将原件输出
        self.zip_out = "{}/{}".format(output_dir, 'zip_out')  # 原件压缩包解压后的存储路径
        self.dcm_out = "{}/{}".format(output_dir, 'dcm_out')  # 修改操作后的dcm文件存储路径
        self.complete = '{}/{}.txt'.format(outs_dir, 'complete')  # 操作完成后的日志记录
        self.abnormal = '{}/{}.txt'.format(outs_dir, 'abnormal')  # 操作异常的日志记录
        self.file_list = os.listdir(self.input_dir)  # 读取压缩包的list列表
        self.setting_configs = dict()  # 定义一个字典
        for nas in desensitization_switch:  # 获取配置字段文件的值
            if desensitization_switch[nas] is True:
                self.setting_configs[nas] = desensitization_tags[nas]
        if os.path.exists(self.output_dir) is False:  # 创建 输出目录
            os.mkdir(self.output_dir)
        if os.path.exists(self.except_dir) is False:  # 创建异常输出目录
            os.mkdir(self.except_dir)
        if os.path.exists(self.complete) is False:  # 创建 操作完成后的日志记录.txt
            open(self.complete, 'w', encoding='utf-8')
        if os.path.exists(self.abnormal) is False:  # 创建 操作异常的日志记录.txt
            open(self.abnormal, 'w', encoding='utf-8')

    def run(self):
        """ 判断input_dir 目录 为压缩包则调用递归解压方法reads_zip(), 否则 直接Copy至 output_dir目录
            解压后调用 reads_dcm() 传入压缩包解压后的路径，进行dcm文件内容读取，修改等操作
        :return:
        """
        for na in self.file_list:
            if na in open(self.complete, encoding='utf-8').read():  # 如果该压缩包名存在操作完成日志文件中则跳过操作
                continue
            paths = os.path.join(self.input_dir, na)
            if zipfile.is_zipfile(paths):
                logger.info(f"Start decompression {paths}")  # 开始解压
                out_path = os.path.join(self.zip_out, str(na).strip('.zip'))
                zip_files = zipfile.ZipFile(paths)
                zip_lis = zip_files.namelist()
                for dcm in zip_lis:
                    zip_files.extract(dcm, out_path)
                zip_files.close()
                except_outs = os.path.join(self.except_dir, na)  # 定义异常输出路径
                fis = open(self.abnormal, 'a', encoding='utf-8')  # 打开异常日志记录文件
                try:
                    self.reads_zip(out_path)  # 处理压缩包函数
                    self.reads_dcm(out_path)  # 调用函数 读取解压完成的dcm文件
                    # 将修改输出的dcm直接压缩,1、压缩文件的文件名，2、解压格式 3、将制定文件夹进行压缩
                    shutil.make_archive(base_name="{}/{}".format(self.output_dir, str(na).strip('.zip')), format='zip', root_dir="{}/{}".format(self.dcm_out, str(na).strip('.zip')))
                    logger.info(f"Operation succeeded {paths}")  # 操作成功 日志
                    fi = open(self.complete, 'a', encoding='utf-8')  # 操作完成日志
                    fi.write('{}\n'.format(na))
                    fi.close()
                except KeyboardInterrupt:
                    shutil.copytree(out_path, except_outs)  # 异常则直接拷贝
                    fis.write("{} - 用户中断操作".format(paths))
                    print('用户中断执行')
                    break
                except StopIteration:
                    shutil.copytree(out_path, except_outs)  # 异常则直接拷贝
                    fi.write("{}：StopIteration-迭代器没有更多的值\n".format(paths))
                except RuntimeError:
                    shutil.copytree(out_path, except_outs)  # 异常则直接拷贝
                    fi.write("{}：RuntimeError-运行时错误\n".format(paths))
                except IOError:
                    shutil.copytree(out_path, except_outs)  # 异常则直接拷贝
                    fi.write("{}：IOError-输入/输出操作失败\n".format(paths))
                except OSError:
                    shutil.copytree(out_path, except_outs)  # 异常则直接拷贝
                    fi.write("{}：OSError-操作系统错误\n".format(paths))
                except Exception:
                    fi.write("{}：OSError-其他异常\n".format(paths))
                    shutil.copytree(out_path, except_outs)  # 异常拷贝
                    continue
                fis.close()  # 关闭异常日志记录文件
            else:
                copy_out = "{}/{}".format(self.output_dir, na)
                if os.path.exists(copy_out) is True:
                    shutil.rmtree(copy_out)
                logger.error(f"{paths} Copy {copy_out}")  # 无效zip 拷贝
                shutil.copytree(paths, copy_out)
        if os.path.exists(self.zip_out) is True:
            shutil.rmtree(self.zip_out)
        if os.path.exists(self.dcm_out) is True:
            shutil.rmtree(self.dcm_out)

    def reads_dcm(self, src_dcm):
        """ src_dcm = 当前解压完的压缩包顶层路径 读取解压完的内容
            当os.walk() f 为dcm文件列表则进入操作
            tqdm(f = list) 创建进度条对象
        :param src_dcm:
        :return:  dest_directory 将修改后输出dcm文件的目录路径返回
        """
        src_len = len(self.zip_out)
        for r, d, f, in list(os.walk(src_dcm)):
            if len(f):
                bar = tqdm(f)  # 进度条显示
                for fns in bar:
                    directory = r[src_len + 1:]  # 根据input字符串长度 获取 input路径后的文件目录名
                    bar.set_description(f"{directory}")
                    dest_directory = r"{}/{}".format(self.dcm_out, directory)  # 输出目录dcm路径 + 当前input后的路径
                    if os.path.exists(dest_directory) is False:
                        os.makedirs(dest_directory)  # 迭代创建
                    src_input = os.path.join(r, fns)  # 读取的 dcm文件路径
                    self.out_dcm(src_input, dest_directory, fns)  # 读取dcm路径， 输出dcm路径，dcm文件名 -- 调用了操作dcm文件img对象操作函数
        return dest_directory

    def out_dcm(self, src_input, dest_directory, fns):
        """ 修改dcm内容 - 保存输出
            调用读取dcm img操作对象 方法
            img 对象异常 则将dcm文件直接拷贝至输出目录
            img对象正常则 将img对象的内容写入到 dest_directory 目录
        """
        img = self.img_operation(src_input)
        if img is None:
            shutil.copy(src_input, dest_directory)
        else:
            img.save_as("{}/{}".format(dest_directory, fns))

    def img_operation(self, src_input):
        """ dcm文件内容读取
            判断img对象读取成功，可进行img对象属性操作
            读取失败返回 None
            setting_configs = {} 配置需要修改的值
        :param src_input:
        :return:
        """
        try:
            img = pydicom.read_file(src_input, force=True)
            logger.info(f'Image reading succeeded {src_input}')
        except RuntimeError:
            logger.error(f'Image reading failed {src_input}')
            return
        img.update(self.setting_configs)  # 此处可做img对象操作
        logger.info(f'Image operation succeeded {src_input}')
        return img

    def reads_zip(self, src):  # 解压-压缩包
        """ 递归调用压缩
            逻辑：如果是压缩包则解压至当前目录，由于os.walk()  会自动遍历每个目录 返回 r, d, f
            直接再原路径的基础上输出
            使用zip_file 库 进行解压至当前文件目录
            解压完，将当前访问路径压缩包删除
        """
        outs = ''
        for rs, ds, fi in list(os.walk(src)):
            if len(fi):
                for fn in fi:
                    if zipfile.is_zipfile("{}/{}".format(rs, fn)):
                        outs = r"{}/{}".format(rs, str(fn).strip('.zip'))
                        zip_files1 = zipfile.ZipFile("{}/{}".format(rs, fn))
                        zip_list = zip_files1.namelist()
                        for dcm in zip_list:
                            zip_files1.extract(dcm, outs)
                        zip_files1.close()
                        logger.info(f'The subdirectory is unzipped {"{}/{}".format(outs, dcm)}')
                        if os.path.exists("{}/{}".format(rs, fn)) is True:
                            os.remove("{}/{}".format(rs, fn))
                    else:
                        pass
                self.reads_zip(src=outs)


def main():
    args = docopt(__doc__, version="0.1.0")
    input_dir = args.get("<input_dir>", None)
    output_dir = args.get("<output_dir>", None)
    except_dir = args.get("<except_dir>", None)
    runs = BatchDeId(input_dir, output_dir, except_dir)
    runs.run()


if __name__ == '__main__':
    """ 影像支持zip输入，zip输出。 -- 完成
        支持脱敏字段配置 -- 开始处理
        处理日志要打印过程。 -- 完成 
        批量支持万级影像的脱敏影像处理
        测试命令： python model.py D:\python_ShiXi\flask_demo\tuomin_upgrade\ceshi_test\input D:\python_ShiXi\flask_demo\tuomin_upgrade\ceshi_test\output D:\python_ShiXi\flask_demo\tuomin_upgrade\ceshi_test
    """
    start = time.time()
    inputs = r'D:\python_ShiXi\flask_demo\tuomin_upgrade\ceshi_test\input'
    outputs = r'D:\python_ShiXi\flask_demo\tuomin_upgrade\ceshi_test\out'
    runs = BatchDeId(inputs, outputs)
    runs.run()
    main()
    print("耗时:", time.time() - start, "s")  # 时间截 的 减法


# 解压思路逻辑：先将最外层的压缩包解压，再传入解压的路径，再调用该函数，os.walk() 会遍历该目录，是否存在压缩包，存在的话读取压缩包的内容，解压至当前目录，再rs的路径基础上，将压缩包名去除后缀扩展名进行保存
# 再递归调用该函数，将解压后的路径再次传入，再次判断是否存在压缩包，如果不存在压缩包则结束，有压缩包的话，会一直递归调用，再rs的路径基础上进行输出，再将解压后的压缩包原件删除
# 支持脱敏字段配置，在本地留一个配置文件.txt， 内容对于着key = 属性，value = ["修改 or 删除 ","修改后的值"]
# 处理日志要打印过程， -- 完成


















